﻿#include "dx_descriptor.h"
#include "dx_exception.h"

namespace ifire::dx {
void DxDescriptor::Init(ID3D12Device* device) {
  device_ = device;

  // 需要得到每个描述符的大小
  RTV_SIZE =
      device_->GetDescriptorHandleIncrementSize(D3D12_DESCRIPTOR_HEAP_TYPE_RTV);
  DSV_SIZE =
      device_->GetDescriptorHandleIncrementSize(D3D12_DESCRIPTOR_HEAP_TYPE_DSV);
  CBV_SRV_UAV_SIZE = device_->GetDescriptorHandleIncrementSize(
      D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV);

  // 创建Render Target View
  BuildRtv();
  // 创建Depth/Stencil View
  BuildDsv();
}
void DxDescriptor::BuildRtv() {
  // 用于计算每个项目的索引
  // 默认的SwapChain需要2个
  mRtvSwapChainIndex = rtv_index_.count(SWAP_CHAIN_BUFFER_COUNT);
  // +1用于SSAO的屏幕空间法线
  mRtvScreenNormalMapIndex = rtv_index_.count(1);
  // +2用于SSAO的环境光贴图
  mRtvAmbientMapIndex = rtv_index_.count(2);
  // +1用于屏幕后处理
  mRtvOffscreenIndex = rtv_index_.count(1);

  // 创建渲染目标视图资源
  D3D12_DESCRIPTOR_HEAP_DESC rtv_desc = {};

  // rtv_index_count_时此就是总计的数量
  rtv_desc.NumDescriptors = rtv_index_.size();
  rtv_desc.Type = D3D12_DESCRIPTOR_HEAP_TYPE_RTV;
  rtv_desc.Flags = D3D12_DESCRIPTOR_HEAP_FLAG_NONE;
  rtv_desc.NodeMask = 0;
  HASSERT(device_->CreateDescriptorHeap(
      &rtv_desc, IID_PPV_ARGS(rtv_heap_.GetAddressOf())));
}

void DxDescriptor::BuildDsv() {
  // 创建深度/模板视图资源
  D3D12_DESCRIPTOR_HEAP_DESC dsv_desc = {};
  const int shadow_count = 1; // 阴影贴图需要一个DSV
  const int ui_count = 1;     // UI需要一个DSV
  dsv_desc.NumDescriptors = 1 + shadow_count + ui_count;
  dsv_desc.Type = D3D12_DESCRIPTOR_HEAP_TYPE_DSV;
  dsv_desc.Flags = D3D12_DESCRIPTOR_HEAP_FLAG_NONE;
  dsv_desc.NodeMask = 0;
  HASSERT(device_->CreateDescriptorHeap(
      &dsv_desc, IID_PPV_ARGS(dsv_heap_.GetAddressOf())));
}

CD3DX12_CPU_DESCRIPTOR_HANDLE DxDescriptor::GetCpuHandle(int offset) {
  return CD3DX12_CPU_DESCRIPTOR_HANDLE(
      srv_heap_->GetCPUDescriptorHandleForHeapStart(), offset,
      CBV_SRV_UAV_SIZE);
  ;
}

CD3DX12_GPU_DESCRIPTOR_HANDLE DxDescriptor::GetGpuHandle(int offset) {
  return CD3DX12_GPU_DESCRIPTOR_HANDLE(
      srv_heap_->GetGPUDescriptorHandleForHeapStart(), offset,
      CBV_SRV_UAV_SIZE);
}

CD3DX12_CPU_DESCRIPTOR_HANDLE DxDescriptor::GetDsv(int offset) {
  return CD3DX12_CPU_DESCRIPTOR_HANDLE(
      dsv_heap_->GetCPUDescriptorHandleForHeapStart(), offset, DSV_SIZE);
}

CD3DX12_CPU_DESCRIPTOR_HANDLE DxDescriptor::GetRtv(int offset) const {
  return CD3DX12_CPU_DESCRIPTOR_HANDLE(
      rtv_heap_->GetCPUDescriptorHandleForHeapStart(), offset, RTV_SIZE);
}

// CD3DX12_CPU_DESCRIPTOR_HANDLE DxDescriptor::GetTextureStart() {
//   texture_handle_ = GetCpuHandle(mTextureIndex);
//   return texture_handle_;
// }

CD3DX12_GPU_DESCRIPTOR_HANDLE DxDescriptor::GetTextureGpu(int texture_index) {
  return GetGpuHandle(texture_index + mTextureStartIndex);
}

void DxDescriptor::BuildDesc() {
  // 用于普通的2D纹理
  mImageDesc = {};
  mImageDesc.Shader4ComponentMapping = D3D12_DEFAULT_SHADER_4_COMPONENT_MAPPING;
  mImageDesc.Format = BACK_BUFFER_FORMAT;
  mImageDesc.ViewDimension = D3D12_SRV_DIMENSION_TEXTURE2D;
  mImageDesc.Texture2D.MipLevels = -1;

  // 用于SkyBox
  mCubeDesc = mImageDesc;
  mCubeDesc.ViewDimension = D3D12_SRV_DIMENSION_TEXTURECUBE;

  // 用于纯空的纹理
  mNullDesc = mImageDesc;
  mNullDesc.ViewDimension = D3D12_SRV_DIMENSION_TEXTURE2D;
  mNullDesc.Format = DXGI_FORMAT_R8G8B8A8_UNORM;
  mNullDesc.Texture2D.MostDetailedMip = 0;
  mNullDesc.Texture2D.MipLevels = 1;
  mNullDesc.Texture2D.ResourceMinLODClamp = 0.0f;

  // 用于纹理数组
  mTexArrayDesc = mImageDesc;
  mTexArrayDesc.ViewDimension = D3D12_SRV_DIMENSION_TEXTURE2DARRAY;
  mTexArrayDesc.Texture2DArray.MostDetailedMip = 0;
  mTexArrayDesc.Texture2DArray.MipLevels = -1;
  mTexArrayDesc.Texture2DArray.FirstArraySlice = 0;
}

void DxDescriptor::BuildSrv(const DxDescriptorDesc& desc) {
  //
  // 计算SRV堆的大小，固定大小的东西在前，非固定的位置在后。
  //
  index_.reset();

  // 空白天空盒纹理，纯空
  mNullCubeIndex = index_.count(1);
  // 空白2D纹理，纯空
  mNullTextureIndex = index_.count(1);
  // 默认法线纹理，不是空的，是一个1024x1024的标准法线
  mEmptyNormalIndex = index_.count(1);
  // 天空盒的纹理
  mSkyBoxIndex = index_.count(1);
  // 阴影的纹理
  mShadowMapIndex = index_.count(1);
  // SSAO的纹理
  mSsaoIndex = index_.count(3);
  // SSAO Ambient的纹理
  mSsaoAmbientIndex = index_.count(2);
  // 用于波浪
  mWavesIndex = index_.count(desc.WavesCount);
  // 用于索贝尔算子
  mSobelIndex = index_.count(desc.SobelCount);
  // 用于离屏纹理
  mOffscreenIndex = index_.count(1);
  // 用于波浪的置换贴图
  mDisplacementIndex = index_.count(1);
  // UI的纹理
  mUIDepthMapIndex = index_.count(1);
  // 其它的纹理
  mTextureStartIndex = index_.count(desc.TextureCount);

  //
  // 创建SRV堆
  // 此时的index_.count正好是所有项目数量的总合，以此创建。
  //
  D3D12_DESCRIPTOR_HEAP_DESC srvHeapDesc = {};
  srvHeapDesc.NumDescriptors = index_.size();
  srvHeapDesc.Type = D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV;
  srvHeapDesc.Flags = D3D12_DESCRIPTOR_HEAP_FLAG_SHADER_VISIBLE;
  HASSERT(device_->CreateDescriptorHeap(&srvHeapDesc, //
      IID_PPV_ARGS(&srv_heap_)));

  // ------
  // 为项目创建Desc，每个类型不需要重新在每次的时候创建（其实也节约不了多少）
  // ------
  BuildDesc();

  // ------
  // 创建空白项目
  // ------
  BuildNullSrv();

  // BuildEmptyNormalMap(desc.EmptyNormal);
  // BuildSkyBox(desc.SkyBox);

  // 纹理、阴影贴图在外面创建
}

void DxDescriptor::BuildNullSrv() {
  // ------
  // 空白的天空盒
  // ------
  auto null_cube = GetCpuHandle(mNullCubeIndex);
  device_->CreateShaderResourceView(nullptr, &mCubeDesc, null_cube);
  // 得到GPU引用
  mNullCube = GetGpuHandle(mNullCubeIndex);

  // ------
  // 空白的纹理
  // ------
  auto null_tex = GetCpuHandle(mNullTextureIndex);
  device_->CreateShaderResourceView(nullptr, &mNullDesc, null_tex);
  // 得到GPU引用
  mNullTexture = GetGpuHandle(mNullTextureIndex);
}

void DxDescriptor::AssignEmptyNormalMap(ID3D12Resource* res) {
  // 根据资源本身的格式来设置视图
  mImageDesc.Format = res->GetDesc().Format;

  auto handle = GetCpuHandle(mEmptyNormalIndex);
  device_->CreateShaderResourceView(res, &mImageDesc, handle);
  // 得到GPU引用
  mEmptyNormal = GetGpuHandle(mEmptyNormalIndex);
}

void DxDescriptor::AssignSkyBox(ID3D12Resource* res) {
  mCubeDesc.Format = res->GetDesc().Format;
  mCubeDesc.TextureCube.MipLevels = res->GetDesc().MipLevels;
  auto handle = GetCpuHandle(mSkyBoxIndex);
  device_->CreateShaderResourceView(res, &mCubeDesc, handle);
  // 得到GPU引用
  mSkyBox = GetGpuHandle(mSkyBoxIndex);
}

UINT DxDescriptor::CreateTexture(TextureType type, ID3D12Resource* res) {
  // 得到Texture的索引的位置
  auto texture_handle_ =
      GetCpuHandle(mTextureStartIndex + mTextureCurrentIndex);

  // 如果是Cube，需要不同的DESC
  if (type == TextureType::Cube) {
    // 根据资源本身的格式来设置视图
    mCubeDesc.Format = res->GetDesc().Format;
    mCubeDesc.TextureCube.MipLevels = res->GetDesc().MipLevels;
    device_->CreateShaderResourceView(res, &mCubeDesc, texture_handle_);
  } else if (type == TextureType::TextureArray) {
    mTexArrayDesc.Format = res->GetDesc().Format;
    mTexArrayDesc.Texture2DArray.ArraySize = res->GetDesc().DepthOrArraySize;
    device_->CreateShaderResourceView(res, &mTexArrayDesc, texture_handle_);
  } else {
    // 根据资源本身的格式来设置视图
    mImageDesc.Format = res->GetDesc().Format;
    device_->CreateShaderResourceView(res, &mImageDesc, texture_handle_);
  }

  // 步进索引，并返回当前的索引
  auto offset_index = mTextureCurrentIndex++;
  return offset_index;
}

} // namespace ifire::dx
